import cv2
import numpy as np
from config import SSIM_SIZE, PHASH_SIZE, ALMOST_BLACK_BITMAP


class FeatureExtractor:
    def __init__(self):
        self._ssim_size = SSIM_SIZE
        self._pHashSize = PHASH_SIZE
        self._almostBlackBitmap = ALMOST_BLACK_BITMAP

    def extract_frame_feature(self, frame):
        """提取帧的特征"""
        features = {}

        features.update(self._extract_basic_visual_features(frame))

        return features

    def _extract_basic_visual_features(self, frame):
        """提取基础视觉特征"""
        # 1. 调整大小
        resized = cv2.resize(frame, (64, 64))

        # 2. 转换为HSV颜色空间
        hsv = cv2.cvtColor(resized, cv2.COLOR_BGR2HSV)

        # 3. 计算颜色直方图
        hist = cv2.calcHist([hsv], [0, 1], None, [8, 8], [0, 180, 0, 256])
        hist = cv2.normalize(hist, hist).flatten()

        # 4. 计算边缘特征
        gray = cv2.cvtColor(resized, cv2.COLOR_BGR2GRAY)
        edges = cv2.Canny(gray, 100, 200)
        edge_density = np.sum(edges) / (edges.shape[0] * edges.shape[1])

        # 5. 计算DCT特征
        dct = cv2.dct(gray.astype(np.float32))
        dct_feature = dct[:8, :8].flatten()

        return {
            'hist': hist,
            'edge_density': edge_density,
            'dct': dct_feature
        }